`include "mycpu.h"

module mem_stage(
    input                          clk           ,
    input                          reset         ,
    input                          flush         ,
    //allowin
    input                          ws_allowin    ,
    output                         ms_allowin    ,
    output                         ms_flush      ,
    output                         ms_inst_mtc0_v,
    //from es
    input                          es_to_ms_valid,
    input  [`ES_TO_MS_BUS_WD -1:0] es_to_ms_bus  ,
    //to ws
    output                         ms_to_ws_valid,
    output [`MS_TO_WS_BUS_WD -1:0] ms_to_ws_bus  ,
    //to ds
    output                         ms_to_ds_valid,
    output [`MS_TO_DS_BUS_WD -1:0] ms_to_ds_bus,
    //to fs
    output                   [31:0]ms_pc,
    //from data-sram
    input         data_sram_data_ok,
    input [31:0]  data_sram_rdata,
    input [63:0]  mulresult
);

reg         ms_valid;
wire        ms_ready_go;
reg   [31:0]save_data;//use to save data for several cycle
reg         save;//use to indicate whether data in save_data is valid
reg         abandon;//whether abandon the next data we received
reg         flush_r;//to save flush_r

reg [`ES_TO_MS_BUS_WD -1:0] es_to_ms_bus_r;
wire        ms_res_from_mem;
wire        ms_gr_we;
wire [ 1:0] ms_addr_4;
reg  [ 3:0] ms_gr_we_4;
wire [ 4:0] ms_dest;
wire [31:0] ms_alu_result;

wire    [31:0] divouthi, divoutlo;
wire  fmul,fdiv;

wire mthi,mtlo,inst_mfhi,inst_mflo;
wire [31:0] hilo;

wire es_ex;
wire es_mem_we;
wire [4:0]ms_excode;
wire [4:0]ms_c0_waddr;
wire  [31:0] ms_badvaddr;
wire ms_inst_tlbwi;
wire ms_inst_tlbr;
wire ms_bd;
wire ms_inst_mfc0;
wire ms_inst_eret;
wire ms_inst_lwl;
wire ms_inst_lwr;
wire ms_load_sign;
wire ms_load_byte;
wire ms_load_hw;
wire ms_store_byte;
wire ms_store_hw;
wire ms_tlb_r;
wire ms_refetch;


assign {ms_refetch ,//231
        ms_tlb_r,   //230
        ms_inst_tlbwi,//229
        ms_inst_tlbr,//228
        es_mem_we,   //227:227
        ms_badvaddr , 
        ms_c0_waddr ,
        ms_bd       ,
        es_ex       ,
        ms_excode   ,
        ms_inst_mfc0,
        ms_inst_mtc0,
        ms_inst_eret,    //194:180
        ms_inst_lwl,
        ms_inst_lwr,
        ms_load_sign,
        ms_load_byte,
        ms_load_hw,
        ms_store_byte,
        ms_store_hw,     //179: 173
        mthi,
        mtlo,
        inst_mfhi,
        inst_mflo,
        hilo, 
        fmul,
        fdiv,
        divouthi,
        divoutlo,         //172:71
        ms_res_from_mem,  //70:70
        ms_gr_we       ,  //69:69
        ms_dest        ,  //68:64
        ms_alu_result  ,  //63:32
        ms_pc             //31:0
       } = es_to_ms_bus_r;

reg  [31:0] mem_result;
wire [31:0] ms_final_result;
wire [3:0]  ms_rf_we;

wire ms_ex;

assign ms_inst_mtc0_v=ms_inst_mtc0&ms_valid;


reg [31:0] LO;
reg [31:0] HI;



assign ms_to_ws_bus = {ms_refetch,    //123
                       ms_tlb_r,     //122
                       ms_inst_tlbwi,//121
                       ms_inst_tlbr,//120
                       ms_badvaddr,  //   119:88
                       ms_c0_waddr,  //87:83
                       ms_bd       ,
                       ms_ex      ,
                       ms_excode   , //80:76
                       ms_inst_mfc0,
                       ms_inst_mtc0,
                       ms_inst_eret,    //73
                       ms_gr_we_4     ,  //72:69
                       ms_dest        ,  //68:64
                       ms_final_result,  //63:32
                       ms_pc             //31:0
                      };
assign ms_to_ds_bus = {ms_ready_go,         //43:43
                       ms_inst_mfc0   ,   //42:42
                       ms_rf_we,         //41:38
                       ms_res_from_mem,   //37:37
                       ms_dest        ,  //36:32
                       ms_final_result    //31:0
                      };
// handle if ws is blocked,but can not happen till current experiment
always @(posedge clk) begin
    if(reset)begin
        save <= 1'b0;
    end
    else if((flush|flush_r)&abandon)begin
        save <= 1'b0;
    end
    else if(ms_ready_go & ~ws_allowin)begin
        save <= 1'b1;
    end
    else begin
        save <= 1'b0;
    end
end
always @(posedge clk) begin
    if(reset)begin
        save_data <= 32'b0;
    end
    else if(ms_ready_go & ~ws_allowin)begin
        save_data <= data_sram_rdata;
    end
    else begin
        save_data <= save_data;
    end
end
always @(posedge clk) begin
    if(reset)begin
        abandon <= 1'b0;
    end
    else if(data_sram_data_ok) begin
        abandon <= 1'b0;
    end
    else if(~ms_ex & ms_valid & ~flush_r & ~flush & ((ms_res_from_mem & ~inst_mfhi & ~inst_mflo)| es_mem_we) & es_to_ms_valid)begin
        abandon <= 1'b1;
    end
    /*else if(~ms_ex & ms_valid & ~flush_r & ~flush & ((ms_res_from_mem & ~inst_mfhi & ~inst_mflo)| es_mem_we) & ~ms_allowin & ~ms_ready_go)begin
        abandon <= 1'b1;
    end*/
    else begin
        abandon <= abandon;
    end
end
always @(posedge clk) begin
    if(reset)begin
        flush_r <= 1'b0;
    end
    else if(es_to_ms_valid & ms_allowin & ~(abandon & ~data_sram_data_ok))begin
        flush_r <= 1'b0;
    end
    else if(flush)begin
        flush_r <= 1'b1;
    end
    else begin
        flush_r <= flush_r;
    end
end

//exception
assign ms_ex=es_ex & ms_valid;
assign ms_flush=ms_ex | (ms_inst_eret&ms_valid);

assign ms_ready_go    = (~(~ms_ex & ms_valid & ((ms_res_from_mem & ~inst_mfhi & ~inst_mflo)| es_mem_we)) | (data_sram_data_ok | save))& ~(abandon&(flush_r|flush));
assign ms_allowin     = !ms_valid || ms_ready_go && ws_allowin;
assign ms_to_ws_valid = ms_valid && ms_ready_go;
assign ms_to_ds_valid = ms_valid ;
assign ms_rf_we      =  ms_gr_we_4;
assign ms_addr_4     =  ms_alu_result[1:0];
always @(*) begin
    if(ms_gr_we)begin
        if(ms_inst_lwl)begin
            ms_gr_we_4 = (ms_addr_4==2'b00)?4'b1000:
                        (ms_addr_4==2'b01)?4'b1100:
                        (ms_addr_4==2'b10)?4'b1110:
                        /*(ms_addr_4==2'b11)?*/4'b1111;
        end
        else if(ms_inst_lwr)begin
            ms_gr_we_4 = (ms_addr_4==2'b00)?4'b1111:
                        (ms_addr_4==2'b01)?4'b0111:
                        (ms_addr_4==2'b10)?4'b0011:
                        /*(ms_addr_4==2'b11)?*/4'b0001;
        end
        else begin
            ms_gr_we_4 = 4'b1111;
        end
    end
    else begin
        ms_gr_we_4 = 4'b0000;
    end
    
end

always @(posedge clk) begin
    if(reset) begin
        LO <= 32'b0;
        HI <= 32'b0;
    end
    else if(!flush && !ms_ex && ms_valid) begin  
        if(fmul)begin
            LO <= mulresult[31:0];
            HI <= mulresult[63:32];
        end
        else if(fdiv) begin
            LO <= divoutlo;
            HI <= divouthi;
        end
        else if(mthi) HI <=  hilo;
        else if(mtlo) LO <=  hilo;
    end
end

always @(posedge clk) begin
    if (reset|flush) begin
        ms_valid <= 1'b0;
    end
    else if (ms_allowin) begin
        ms_valid <= es_to_ms_valid;
    end
    else if (flush_r) begin
        ms_valid <= 1'b0;
    end
    if (es_to_ms_valid && ms_allowin) begin
        es_to_ms_bus_r  <= es_to_ms_bus; 
    end
end

always @(*)begin
    if(ms_load_byte) begin
        if(ms_load_sign)begin
            mem_result = (ms_addr_4==2'b00)?{{24{data_sram_rdata[7]}},data_sram_rdata[7:0]}:
                        (ms_addr_4==2'b01)?{{24{data_sram_rdata[15]}},data_sram_rdata[15:8]}:
                        (ms_addr_4==2'b10)?{{24{data_sram_rdata[23]}},data_sram_rdata[23:16]}:
                        /*(ms_addr_4==2'b11)*/{{24{data_sram_rdata[31]}},data_sram_rdata[31:24]};
        end
        else begin
            mem_result = (ms_addr_4==2'b00)?{24'b0,data_sram_rdata[7:0]}:
                        (ms_addr_4==2'b01)?{24'b0,data_sram_rdata[15:8]}:
                        (ms_addr_4==2'b10)?{24'b0,data_sram_rdata[23:16]}:
                        /*(ms_addr_4==2'b11)*/{24'b0,data_sram_rdata[31:24]};
        end
    end
    else if(ms_load_hw) begin
        if(ms_load_sign)begin
            mem_result = (ms_addr_4==2'b00)?{{16{data_sram_rdata[15]}},data_sram_rdata[15:0]}:
                        /*(ms_addr_4==2'b10)?*/{{16{data_sram_rdata[31]}},data_sram_rdata[31:16]};
        end
        else begin
            mem_result = (ms_addr_4==2'b00)?{16'b0,data_sram_rdata[15:0]}:
                        /*(ms_addr_4==2'b10)?*/{16'b0,data_sram_rdata[31:16]};
        end
    end
    else if(ms_inst_lwl) begin
        mem_result =    (ms_addr_4==2'b00)?{data_sram_rdata[7:0],24'b0}:
                        (ms_addr_4==2'b01)?{data_sram_rdata[15:0],16'b0}:
                        (ms_addr_4==2'b10)?{data_sram_rdata[23:0],8'b0}:
                        /*(ms_addr_4==2'b11)*/{data_sram_rdata[31:0]};
    end
    else if(ms_inst_lwr) begin
        mem_result =    (ms_addr_4==2'b00)?{data_sram_rdata[31:0]}:
                        (ms_addr_4==2'b01)?{8'b0,data_sram_rdata[31:8]}:
                        (ms_addr_4==2'b10)?{16'b0,data_sram_rdata[31:16]}:
                        /*(ms_addr_4==2'b11)*/{24'b0,data_sram_rdata[31:24]};
    end
    else begin //load word
        mem_result = data_sram_rdata;
    end
end

assign ms_final_result = (inst_mfhi)  ?  HI :
                         (inst_mflo) ?  LO :
                         ms_res_from_mem ? ((save)?save_data:mem_result ):
                         ms_alu_result;

endmodule
